# -*- coding: utf-8 -*-
"""
利用感知机进行文本分类
"""

from sklearn.datasets import fetch_20newsgroups
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import Perceptron
from sklearn.metrics import classification_report
from sklearn.externals import joblib

# sklearn内置数据集，该数据集包含了来自 20 个 Usenet 新闻组约 20000 份文档样本。
categories = ['rec.sport.hockey', 'rec.sport.baseball','rec.autos']
newsgroups_train = fetch_20newsgroups(subset='train',categories=categories, remove=('headers', 'footers', 'quotes'))
newsgroups_test = fetch_20newsgroups(subset='test',categories=categories, remove=('headers', 'footers', 'quotes'))

tfidf_vec = TfidfVectorizer()
tfidf_vec.fit(newsgroups_train.data)
joblib.dump(tfidf_vec,'tfidf_vec.model')  # 保存模型

X_train = tfidf_vec.transform(newsgroups_train.data)
X_test = tfidf_vec.transform(newsgroups_test.data)

clf = Perceptron(random_state=11)
clf.fit(X_train, newsgroups_train.target)
joblib.dump(clf, 'clf.model')  # 保存模型

predictions = clf.predict(X_test)
print(classification_report(newsgroups_test.target, predictions))


